from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
from numpy import sin, cos, mgrid, pi, sqrt


class Form(HasTraits):
    engine = Instance(Engine, ())
    scene = Instance(MlabSceneModel)

    def _scene_default(self):
        self.engine.start()
        scene = MlabSceneModel(engine=self.engine)
        return scene

    @observe('scene.activated')
    def update_plot(self, event=None):
        u, v = mgrid[- 0.035:pi:0.01, - 0.05:pi:0.01]
        X = 2 / 3. * (cos(u) * cos(2 * v) + sqrt(2) * sin(u) * cos(v)) * cos(u) / (sqrt(2) - sin(2 * u) * sin(3 * v))
        Y = 2 / 3. * (cos(u) * sin(2 * v) - sqrt(2) * sin(u) * sin(v)) * cos(u) / (sqrt(2) - sin(2 * u) * sin(3 * v))
        Z = -sqrt(2) * cos(u) * cos(u) / (sqrt(2) - sin(2 * u) * sin(3 * v))
        S = sin(u)
        self.scene.mlab.mesh(X, Y, Z, scalars=S, colormap='YlGnBu', )
        # self.scene.mlab.view(.0, - 5.0, 4)

    view = View(Item('scene', editor=SceneEditor(scene_class=MayaviScene),
                     height=250, width=300, show_label=False),
                resizable=True  # We need this to resize with the parent widget
                )


from PyQt5.QtWidgets import *
class FormWidget(QWidget):
    def __init__(self):
        QWidget.__init__(self)
        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        self.form = Form()
        layout.addWidget(self.form.edit_traits(parent=self, kind='subpanel').control)
formWidget = FormWidget()
if __name__ == '__main__':
    import os
    import sys

    app = QApplication(sys.argv)
    formWidget.show()
    sys.exit(app.exec_())
